import pandas as pd
from sklearn.tree import DecisionTreeClassifier, export_graphviz
import graphviz

# 定义西瓜数据集
data = pd.DataFrame({
    '编号': [1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17],
    '色泽': ['青绿','乌黑','乌黑','青绿','浅白','青绿','乌黑','乌黑','乌黑','青绿','浅白','浅白','青绿','浅白','乌黑','浅白','青绿'],
    '根蒂': ['蜷缩','蜷缩','蜷缩','蜷缩','蜷缩','稍蜷','稍蜷','稍蜷','稍蜷','硬挺','硬挺','蜷缩','稍蜷','稍蜷','稍蜷','蜷缩','蜷缩'],
    '敲声': ['浊响','沉闷','浊响','沉闷','浊响','浊响','浊响','浊响','沉闷','清脆','清脆','浊响','浊响','沉闷','浊响','浊响','沉闷'],
    '纹理': ['清晰','清晰','清晰','清晰','清晰','稍糊','稍糊','清晰','稍糊','清晰','模糊','模糊','稍糊','稍糊','清晰','模糊','稍糊'],
    '脐部': ['凹陷','凹陷','凹陷','凹陷','凹陷','稍凹','稍凹','稍凹','稍凹','平坦','平坦','平坦','凹陷','凹陷','稍凹','平坦','稍凹'],
    '触感': ['硬滑','硬滑','硬滑','硬滑','硬滑','软粘','软粘','硬滑','硬滑','软粘','硬滑','软粘','硬滑','硬滑','软粘','硬滑','硬滑'],
    '好瓜': ['是','是','是','是','是','是','是','是','否','否','否','否','否','否','否','否','否']
})

# 特征编码
from sklearn.preprocessing import OrdinalEncoder
encoder = OrdinalEncoder()
cat_cols = ['色泽','根蒂','敲声','纹理','脐部','触感','好瓜']
data[cat_cols] = encoder.fit_transform(data[cat_cols])

# 划分特征和目标
X = data.drop(['编号','好瓜'], axis=1)
y = data['好瓜']

# 构建决策树
clf = DecisionTreeClassifier()
clf.fit(X, y)

# 生成可视化数据
dot_data = export_graphviz(
    clf,
    out_file=None,
    feature_names=X.columns,
    class_names=['否','是'],
    filled=True,
    rounded=True,
    special_characters=True
)

# 关键：手动指定 dot 可执行文件路径（替换为你的实际路径）
graph = graphviz.Source(dot_data, engine='dot', executable=r'C:\Program Files\Graphviz\bin\dot.exe')
# 保存为图片（无需调用 render，直接显示或保存）
graph.view(filename='watermelon_decision_tree', cleanup=True)  # 会自动打开图片